from __future__ import annotations

import os
import sys
import datetime
import json
from typing import List, Dict, Optional, Type, Generator, Union, Tuple

import peewee
from peewee import Model, Field, CharField, DateTimeField, PostgresqlDatabase
from playhouse.postgres_ext import BinaryJSONField

from playhouse.shortcuts import model_to_dict

from shapely import wkb
from shapely.geometry.base import BaseGeometry

from common.logger import logger
from common.datetime_utils.timezone_tools import utcnow
from common.geometry.geom import Geom

from catalog.models.config import (
    PG_DATABASE_NAME,
    PG_HOST,
    PG_PORT,
)

# The nas account loads PostGIS password differently.
if os.getuid() == 888187:
    # NOTE The password file is protected, and only the nas account can read it.
    sys.path.insert(0, "/NAS/STAC/CATALOG/PostGIS_OWNER")
    from postgis_owner_config import PG_USER_NAME, PG_PASSWORD
else:
    from catalog.models.config import PG_USER_NAME, PG_PASSWORD


class GeometryField(Field):
    field_type = "geometry"

    def db_value(self, value):
        if isinstance(value, Geom):
            value = value.to_geometry_json()
        elif isinstance(value, BaseGeometry):
            value = json.dumps(value.__geo_interface__)
        return peewee.fn.ST_GeomFromGeoJSON(value)

    def python_value(self, value):
        return Geom(wkb.loads(value, hex=True))


def init_db():
    db = PostgresqlDatabase(
        PG_DATABASE_NAME,
        user=PG_USER_NAME,
        password=PG_PASSWORD,
        host=PG_HOST,
        port=PG_PORT,
        autoconnect=True,
    )
    return db


DEFAULT_DB = init_db()


class BaseItem(Model):
    """Base class for all items managed by the catalog.

    Attributes
    ----------
    uid: CharField
        Unique identifier for this item, primary key.
    valid_until: DateTimeField
        Date and time after which this item is no longer valid, indexed.
    geom: GeometryField
        Spatial geometry for this item, indexed.
    info: BinaryJSONField
        Full information about this item as a json object.
    """

    uid = CharField(max_length=100, primary_key=True)

    valid_until = DateTimeField(null=True)
    geom = GeometryField(index=True, index_type="GIST")
    info = BinaryJSONField(index=False)

    class Meta:
        database = DEFAULT_DB

    @classmethod
    def init_connection(cls: Type[BaseItem]):
        """Initialize a new connection."""
        cls._meta.database.connect(reuse_if_open=True)

    @classmethod
    def close_connection(cls: Type[BaseItem]):
        """
        Close current and idle connections.
        Here we still keep the close operations. This is because the PoolDatabase doesn't allow
        connection exists before multithread / multiprocess starts. So if someone init a connection
        before the multiprocessing, and never close it. It will raise Exception.
        """
        cls._meta.database.close()

    def auto_connection_value(func):
        def connection_wrapper(cls, *args, **kwargs):
            cls.init_connection()
            result = func(cls, *args, **kwargs)
            cls.close_connection()
            return result

        return connection_wrapper

    def auto_connection_generator(func):
        def connection_wrapper(cls, *args, **kwargs):
            cls.init_connection()
            result = func(cls, *args, **kwargs)
            assert isinstance(result, Generator)
            for val in result:
                yield val

            cls.close_connection()

        return connection_wrapper

    @classmethod
    def create_one_item(cls: Type[BaseItem], info_item: Dict) -> BaseItem:
        raise NotImplementedError("Must be provided by a subclass")

    @classmethod
    def create_many_items(
        cls: Type[BaseItem], lst_info_items: List[Dict]
    ) -> Generator[BaseItem, None, None]:
        """Create many items from given list of dictionaries.

        Parameters
        ----------
        lst_info_items : List[Dict]
            A list of dictionary, with each dictionary containing the information of one item.

        Yields
        ------
        Generator[BaseItem, None, None]
            A generator of created items.
        """
        return (cls.create_one_item(info_item) for info_item in lst_info_items)

    @classmethod
    def insert_one_item(
        cls: Type[BaseItem],
        item: Union[BaseItem, Dict],
        on_conflict: str = "ignore",
    ) -> str:
        """Insert one item into the model's table in the database.

        Parameters
        ----------
        item : Union[BaseItem, Dict]
            A dictionary containing the information of one item, or an unsaved item object.
        on_conflict : str, optional
            Strategy used in handling insertion conflict, by default "ignore".
            Currently we support `ignore` (ignore the operation, the original item in the database is kept),
            `update` (update existing item in database),
            and `raise` (raise an error).

        Returns
        -------
        str
            The uid of the inserted item.

        Raises
        ------
        ValueError
            Raise when an unknown `on_conflict` is given.
            Notice that we plan to support `update-unless-invalid` in the future.
        """
        if isinstance(item, dict):
            item = cls.create_one_item(item)

        if on_conflict == "ignore":
            cls.insert(**model_to_dict(item)).on_conflict_ignore().execute()
        elif on_conflict == "update":
            # Postgresql's `on conflict update` is quite different from the typical upsert semantics
            # so let us just do a real delete then insert
            try:
                with cls._meta.database.atomic():
                    cls.insert(**model_to_dict(item)).execute()
            except peewee.IntegrityError as err:
                if "already exists" in str(err):
                    with cls._meta.database.atomic():
                        cls.delete().where(cls.uid == item.uid).execute()
                        cls.insert(**model_to_dict(item)).execute()
                else:
                    logger.error(f"Failed to insert item `{item.uid}` due to `{err}`")
                    raise
        elif on_conflict == "update-unless-invalid":
            # TODO later
            raise NotImplementedError
        elif on_conflict == "raise":
            with cls._meta.database.atomic():
                cls.insert(**model_to_dict(item)).execute()
        else:
            raise ValueError(
                f"Unknown `on_conflict` {on_conflict}, only support `ignore`, `update` and `raise`"
            )
        return item.uid

    @classmethod
    def insert_many_items(
        cls: Type[BaseItem],
        lst_items: List[Dict],
        batch_size: int = 100,
        on_conflict: str = "ignore",
    ):
        """Insert many items into the table.

        Parameters
        ----------
        lst_items : List[Dict]
            A list of dictionary, with each dictionary containing the information of one item.
        batch_size : int, optional
            Batch size for executing the insertion transaction, by default 100
        on_conflict : str, optional
            Strategy used in handling insertion conflict, by default "ignore".
            Currently we support `ignore` (ignore the operation, the original item in the database is kept),
            `update` (update existing item in database),
            and `raise` (raise an error).

        Raises
        ------
        ValueError
            Raise when an unknown `on_conflict` is given.
            Notice that we plan to support `update-unless-invalid` in the future.
        """
        gen_items = (cls.create_one_item(info_item) for info_item in lst_items)
        if on_conflict == "ignore":
            with cls._meta.database.atomic():
                for batch in peewee.chunked(gen_items, batch_size):
                    batch = (model_to_dict(item) for item in batch)
                    cls.insert_many(batch).on_conflict_ignore().execute()
        elif on_conflict == "update":
            # TODO currently just call insert_one_item, but not very efficient
            # might need to change to a customized sql
            for item in gen_items:
                cls.insert_one_item(item, on_conflict="update")
        elif on_conflict == "update-unless-invalid":
            # TODO later
            raise NotImplementedError
        elif on_conflict == "raise":
            with cls._meta.database.atomic():
                for batch in peewee.chunked(gen_items, batch_size):
                    batch = (model_to_dict(item) for item in batch)
                    cls.insert_many(batch).execute()
        else:
            raise ValueError(
                f"Unknown `on_conflict` {on_conflict}, only support `ignore`, `update` and `raise`"
            )

    @classmethod
    def __get_one_item(
        cls: Type[BaseItem],
        uid: str,
        valid_until: Optional[datetime.datetime] = None,
    ) -> Optional[BaseItem]:
        """Get one item from the database by given uid.

        Parameters
        ----------
        uid : str
            The unique id of the item.
        valid_until : Optional[datetime.datetime], optional
            Retrieve item if it is valid at the given datetime,
            by default None to retrieve item that is valid at the time of calling this method.

        Returns
        -------
        Optional[BaseItem]
            None if no item is found, otherwise the item retrieved.
        """
        try:
            item = cls.get_by_id(uid)
            valid_until = utcnow() if valid_until is None else valid_until
            if (item.valid_until is None) or (
                item.valid_until > valid_until.replace(tzinfo=None)
            ):
                return item
            else:
                return None
        except peewee.DoesNotExist:
            return None

    @classmethod
    @auto_connection_value
    def get_one_item(
        cls: Type[BaseItem],
        uid: str,
        valid_until: Optional[datetime.datetime] = None,
    ) -> Optional[BaseItem]:
        """Get one item from the database by given uid.

        Parameters
        ----------
        uid : str
            The unique id of the item.
        valid_until : Optional[datetime.datetime], optional
            Retrieve item if it is valid at the given datetime,
            by default None to retrieve item that is valid at the time of calling this method.

        Returns
        -------
        Optional[BaseItem]
            None if no item is found, otherwise the item retrieved.

        Note
        ----
        Different from `__get_one_item`, this one will auto connect and close the db.
        """
        return cls.__get_one_item(uid, valid_until=valid_until)

    @classmethod
    @auto_connection_generator
    def get_many_items(
        cls: Type[BaseItem],
        uids: List[str],
        valid_until: Optional[datetime.datetime] = None,
    ) -> Generator[Optional[BaseItem], None, None]:
        """Get many items from given list of uids.

        Parameters
        ----------
        uids : List[str]
            A list of unique ids for the items to be retrieved.
        valid_until : Optional[datetime.datetime], optional
            Retrieve items that are valid at the given datetime,
            by default None to only get valid items at the time of calling this method.

        Yields
        ------
        Generator[Optional[BaseItem], None, None]
            A generator of retrieved items.
        """
        # TODO, change to use a more efficient batch way
        for uid in uids:
            yield cls.__get_one_item(uid, valid_until=valid_until)

    @classmethod
    def update_one_item(cls: Type[BaseItem], item: BaseItem):
        """Update one item into table in database."""
        item.save()

    @classmethod
    def update_many_items(cls: Type[BaseItem], items: List[BaseItem]):
        """Update many items into table in database."""
        # TODO, use a sql to do this in batch for efficiency
        for item in items:
            item.save()

    @classmethod
    def delete_one_item(cls: Type[BaseItem], uid: str) -> bool:
        """Delete one item from the model's table in database.

        Parameters
        ----------
        uid : str
            The unique id for the item to be deleted.
            NOTE: we are actually mark the item to be invalidated
            by setting the `valid_until` timestamp for those items, instead of actually delete those items.
            In order to actually delete it, call `sweep_delete_items` after this.

        Returns
        -------
        bool
            Whether the item is deleted successfully.
        """
        try:
            with cls._meta.database.atomic():
                item = cls.__get_one_item(uid)
                if item is not None:
                    query = cls.update(valid_until=datetime.datetime.utcnow()).where(
                        cls.uid == uid
                    )
                    query.execute()
        except Exception as err:
            logger.warn(f"fail to delete {uid} because of {str(err)}")
            return False
        else:
            return True

    @classmethod
    def delete_many_items(cls: Type[BaseItem], uids: List[str]) -> bool:
        """Delete many items from model's table in database.

        Parameters
        ----------
        uids : List[str]
            A list of unique ids for the items to be deleted.
            NOTE: we are actually marking those items to be invalidated
            by setting the `valid_until` timestamp for those items, instead of actually delete those items.
            In order to actually delete them, call `sweep_delete_items` after this.

        Returns
        -------
        bool
            Whether the deletion has succeeded.
        """
        # not actually deleting, just set the valid_until timestamp
        valid_until = datetime.datetime.utcnow()
        try:
            with cls._meta.db.atomic():
                query = cls.update(valid_until=valid_until).where(cls.uid.in_(uids))
                query.execute()
        except Exception as err:
            logger.warn(f"fail to delete {len(uids)} items because of {str(err)}")
            return False
        else:
            return True

    @classmethod
    def delete_many_items_by_query(cls: Type[BaseItem], **kwargs):
        """Delete many items by query with given criterion."""
        # TODO, change to a more efficient impelmentation of using sql's own query to set
        # NOTE, relies on the `query_many_items` provided by a sub-class
        uids = (item.uid for item in cls.query_many_items(**kwargs))
        cls.delete_many_items(uids)

    @classmethod
    def sweep_deleted_items(
        cls: Type[BaseItem],
        valid_before: Optional[datetime.datetime] = None,
        dry_run: bool = True,
    ) -> Union[int, List[BaseItem]]:
        """Sweep deleted items from the database

        Parameters
        ----------
        valid_before : Optional[datetime.datetime], optional
            A timestamp represented by a `datetime.datetime` object,
            all items marked as invalidated before that datetime will be removed from the table,
            by default None to remove all invalidated items.
        dry_run : bool, optional
            Whether to be in the dry-run mode, by default True

        Returns
        -------
        Union[int, List[BaseItem]]
            The number of items removed when `dry_run` is False,
            or a list of items that should be removed when `dry_run` is True.
        """
        valid_before = utcnow() if valid_before is None else valid_before
        if dry_run:
            query = cls.select().where(cls.valid_until < valid_before)
            res = query.execute()
            return list(res)
        else:
            query = cls.delete().where(cls.valid_until < valid_before)
            n = query.execute()
            return n

    @classmethod
    def construct_query(
        cls: Type[BaseItem],
        **kwargs,
    ) -> Tuple[peewee.Query, int]:
        """Construct a query from given criteria,
        return a query object and a total number of criterion to be applied,
        need to be implemented by sub-class"""
        raise NotImplementedError("Must be provided by a subclass")

    @classmethod
    @auto_connection_value
    def count(
        cls: Type[BaseItem],
        spatial: Optional[Geom] = None,
        include_invalidated: bool = False,
        **kwargs,
    ) -> int:
        """Count number of items for this model stored in the database.

        Parameters
        ----------
        spatial : Optional[Geom], optional
            A `Geom` representing the area of interest, by default None.
        include_invalidated : bool, optional
            Whether to include already invalidated items, by default False not to include.

        Returns
        -------
        int
           Number of items matching the given criteria.
        """
        query, _ = cls.construct_query(
            spatial=spatial,
            include_invalidated=include_invalidated,
            **kwargs,
        )
        count = query.count()

        return count

    @classmethod
    def make_table(cls: Type[BaseItem], dry_run: bool = True):
        """Make the table for the model.

        Parameters
        ----------
        dry_run : bool, optional
            Whether to be in dry-run mode, by default True
        """
        logger.info(
            "Create table {} for collection {}, dry-run mode {}".format(
                cls._meta.table_name,
                cls.collection_id,
                dry_run,
            )
        )
        if not dry_run:
            cls.create_table()
            logger.info("Finish creating table {}".format(cls._meta.table_name))

    @classmethod
    def remove_table(cls: Type[BaseItem], dry_run: bool = True):
        """Remove the table of this model.

        Parameters
        ----------
        dry_run : bool, optional
            Whether to be in dry-run mode, by default True
        """
        logger.info(
            "Drop table {} for collection {}, dry-run mode {}".format(
                cls._meta.table_name,
                cls.collection_id,
                dry_run,
            )
        )
        if not dry_run:
            cls._meta.database.drop_tables([cls])
            logger.info("Finish removing table {}".format(cls._meta.table_name))
